import gym
from pprint import pprint

# env = gym.make("highway-v0")
# env.configure({"controlled_vehicles": 2})  # Two controlled vehicles
# env.configure({"vehicles_count": 1})  # A single other vehicle, for the sake of visualisation
# print(env.observation_space)
# print(env.action_space)
# env.reset()
# done = truncated = False
# while not (done or truncated):
#     action = env.action_space.sample()
#     print(action)
#     obs, reward, done, truncated, info = env.step(action)
#     pprint(obs)
#     pprint(reward)
#     pprint(info)
#
#     env.render()



env = gym.make("merge-multi-agent-v0")
print(env.observation_space)
print(env.action_space)
env.reset()
done = truncated = False
while not (done or truncated):
    action = env.action_space.sample()
    pprint(action)
    obs, reward, done, truncated, info = env.step(action)
    pprint(obs)
    pprint(reward)
    pprint(info)
    env.render(mode="rgb_array")
    # from matplotlib import pyplot as plt
    # plt.imshow(env.render(mode="rgb_array"))
    # plt.title("Controlled vehicles are in green")
    # plt.show()


# env = gym.make("highway-multi-agent-v0")
# print("observation_space: ", env.observation_space)
# print("action_space: ", env.action_space)
# env.reset()
# done = truncated = False
# while not (done or truncated):
#     action = env.action_space.sample()
#     print(action)
#     obs, reward, done, truncated, info = env.step(action)
#     pprint(obs)
#     pprint(reward)
#     pprint(info)
#     # env.render(mode="rgb_array")
#     from matplotlib import pyplot as plt
#     plt.imshow(env.render(mode="rgb_array"))
#     plt.title("Controlled vehicles are in green")
#     plt.show()


